import numpy as np
from skimage import measure
from common.img_utils.sas_python.sas_paras import SASPara
from common.img_utils.sas_python.build_bipartite_graph import build_bipartite_graph
from common.img_utils.sas_python.Tcut import Tcut
from common.img_utils.sas_python.calculate_superpixel_props import (
    calculate_superpixel_props,
)


def sas_generator(
    img_arr, sp_labels_list, SAS_paras_type=None, img_region_binary_mask=None
):
    img_arr_height, img_arr_width, _ = img_arr.shape

    # Set SAS parameters for bipartite graph
    para = SASPara(SAS_paras_type)
    if img_region_binary_mask is not None:
        Nseg = int(
            para.Nseg
            * (img_arr_height / 512)
            * (img_arr_width / 512)
            * (
                np.sum(img_region_binary_mask)
                / img_arr_height
                / img_arr_width
                * para.Nseg_offset_a
                + para.Nseg_offset_b
            )
        )
    else:
        Nseg = para.Nseg

    # Construct the post-precessing of superpixels
    seg_all = {}
    seg_vals_all = {}
    seg_edges_all = {}
    for i, sp_label in enumerate(sp_labels_list):
        seg, seg_vals, seg_edges, _ = calculate_superpixel_props(img_arr, sp_label)

        seg_all[str(i)] = seg
        seg_vals_all[str(i)] = seg_vals
        seg_edges_all[str(i)] = seg_edges

    # Perform SAS: Build Bipartite Graph
    print("Build Bipartite Graph...")
    B = build_bipartite_graph(
        [img_arr_height, img_arr_width], para, seg_all, seg_vals_all, seg_edges_all
    )

    # Perform SAS: Transfer Cut
    Nx, Ny = B.shape
    if not Ny < Nseg:
        print("Process Transfer Cut...")
        try:
            label_img = Tcut(B, Nseg, [img_arr_height, img_arr_width])
        except BaseException:
            # If Tcut doesn't work, directly use the last one of sp_labels_list as the final aggregation result
            label_img = sp_labels_list[-1]
    else:
        # If Ny is less than Nseg, directly use the last one of sp_labels_list as the final aggregation result
        label_img = sp_labels_list[-1]

    if not np.any(label_img) > 0:
        # SAS doesn't generate anything
        if img_region_binary_mask is not None:
            label_img += img_region_binary_mask
            label_img = measure.label(label_img, background=0, connectivity=1)

    return label_img
